Two-dimensional derivative information emulators for an SIRS disease model. The code here is an extension of that in the Uncertainty Quantification IV Computer Practicals.
We define the following SIRS disease model and use a numerical solver to solve the differential equations for S, I and R:
# deSolve is used to numerically solve ODEs
library(deSolve)
# Define the SIRSS Disease model. Note x1 = alpha_SI, x2 = alpha_IR, x3 = alpha_SR .
# S = number of susceptible people, I = number of infected people, R = number of recovered people.
SIRS_Disease_Model <- function(t, y, parms) {
with(as.list(parms),{
N <- y["S"] + y["I"] + y["R"]
dS = x3 * y["R"] - x1 * y["S"] * y["I"] / N
dI = x1 * y["S"] * y["I"] / N - x2 * y["I"]
dR = x2 * y["I"] - x3 * y["R"]
res <- c(dS, dI, dR)
list(res)
})
}
We set some initial parameters (midrange values), an initial configuration and a time period up from t=0 to t=100:
# Set input parameters to their midrange values
parms = c( x1 = 0.45, x2 = 0.25, x3 = 0.025 )
# Initial configuration for the number of individuals in compartment S, I and R at time t=0
ystart <- c(S = 850, I = 150, R = 0)
# Define the time period
times <- seq(0, 100, length=1001)
Solve the differential equations and plot the output:
# Run the lsoda solver to solve the differential equations in the SIRS model
out <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000))
# "out" has first column time points, 2nd S, 3rd I and 4th R number
head(out)
## time S I R
## [1,] 0.0 850.0000 150.0000 0.000000
## [2,] 0.1 844.2488 151.9811 3.770086
## [3,] 0.2 838.4716 153.9484 7.580056
## [4,] 0.3 832.6700 155.9005 11.429445
## [5,] 0.4 826.8461 157.8361 15.317763
## [6,] 0.5 821.0017 159.7538 19.244481
# Plot results of the model output
plot_run <- function(){
plot(out[,"time"],out[,"S"],lwd=6,col=4,ty="l",ylim=c(0,max(out[,-1])),xlab="time (t)",ylab="Number of People")
lines(out[,"time"],out[,"I"],lwd=6,col=2)
lines(out[,"time"],out[,"R"],lwd=6,col=3)
legend("topright",legend=c("S = Susceptible"," I = Infected","R = Recovered"),col=c(4,2,3),lwd=6)
}
plot_run()
Bayes linear emulator without using any partial derivative information:
simple_BL_emulator_v2 <- function(x, # the emulator prediction point
xD, # the run input locations xD
D, # the run outputs D = (f(x^1),...,f(x^n))
theta = 1, # the correlation lengths
sigma = 1, # the prior SD sigma sqrt(Var[f(x)])
E_f = 0 # prior expectation of f: E(f(x)) = 0
){
# store length of runs D
n <- length(D)
# Define Covariance structure of f(x): Cov[f(x),f(xdash)]
Cov_fx_fxdash <- function(x,xdash) sigma^2 * exp(-sum((x-xdash)^2)/theta^2)
# Define objects needed for BL update rules
# Create E[D] vector
E_D <- rep(E_f,n)
# Create Var_D matrix:
Var_D <- matrix(0,nrow=n,ncol=n)
for(i in 1:n) for(j in 1:n) Var_D[i,j] <- Cov_fx_fxdash(xD[i,],xD[j,])
# Create E[f(x)]
E_fx <- E_f
# Create Var_f(x)
Var_fx <- sigma^2
# Create Cov_fx_D row vector
Cov_fx_D <- matrix(0,nrow=1,ncol=n)
for(j in 1:n) Cov_fx_D[1,j] <- Cov_fx_fxdash(x,xD[j,])
# Perform Bayes Linear adjustment to find Adjusted Expectation and Variance of f(x)
ED_fx <- E_fx + Cov_fx_D %*% solve(Var_D) %*% (D - E_D)
VarD_fx <- Var_fx - Cov_fx_D %*% solve(Var_D) %*% t(Cov_fx_D)
# Return the emulator adjusted expectation and variance
return(c("ExpD_f(x)"=ED_fx,"VarD_f(x)"=VarD_fx))
}
Create a 16-point maximin Latin hypercube design:
lhd_maximin <- function(nl){ # nl = number of points in LHD
x_lhd <- cbind("x1"=sample(0:(nl-1)),"x2"=sample(0:(nl-1))) / nl + 0.5/nl # create LHD
### Maximin loop: performs swaps on 1st of two closest points with another random point
for(i in 1:1000){
mat <- as.matrix(dist(x_lhd)) + diag(10,nl) # creates matrix of distances between points
# note the inflated diagonal
closest_runs <- which(mat==min(mat),arr.ind=TRUE) # finds pairs of closest runs
ind <- closest_runs[sample(nrow(closest_runs),1),1] # chooses one of close runs at random
swap_ind <- sample(setdiff(1:nl,ind),1) # randomly selects another run to swap with
x_lhd2 <- x_lhd # creates second version of LHD
x_lhd2[ind[1],1] <- x_lhd[swap_ind,1] # swaps x_1 vals between 1st close run & other run
x_lhd2[swap_ind,1] <- x_lhd[ind[1],1] # swaps x_1 vals between 1st close run & other run
if(min(dist(x_lhd2)) >= min(dist(x_lhd))-0.00001) { # if min distance between points is same or better
x_lhd <- x_lhd2 # we replace LHD with new LHD with the swap
# cat("min dist =",min(dist(x_lhd)),"Iteration = ",i,"\n") # write out min dist
}
}
return(x_lhd)
}
set.seed(15)
x_lhd<- lhd_maximin(16)
xD <- x_lhd
### plot maximin Latin hypercube design
plot(xD,xlim=c(0,1),ylim=c(0,1),pch=16,xaxs="i",yaxs="i",col="blue",xlab="x1",ylab="x2",cex=1.4)
abline(h=(0:16)/16,col="grey60")
abline(v=(0:16)/16,col="grey60")
We want to emulate over the 2-dimensional input space of \(x_1\) and \(x_2\) at time t=10, keeping \(x_3=0.04\). Here \(x_1 \in [0.1,0.8]\) and \(x_2 \in [0,0.5]\) so we scale these input ranges to [0,1] for our emulation.
xD_scaled <- cbind("x1"=rep(0,16),"x2"=rep(0,16))
xD_scaled[,1] <- xD[,1]*0.7+0.1
xD_scaled[,2] <- xD[,2]*0.5
We evaluate the true SIRS model for our 16 design runs and extract the model output at t=10 (this corresponds to the 101st row of the output matrix):
# Perform 16 runs of the SIRS model extracting the output for t=10 and store as D
D <- NULL
for(i in 1:nrow(xD_scaled)){
parms = c( xD_scaled[i,1], xD_scaled[i,2], x3 = 0.04)
# Extract the output at t=10, this corresponds to the 101st row
infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000)) [101,"I"]
D <- c(D,infected)
}
Define 50x50 grid of prediction points xP for input into our emulator:
x_grid <- seq(-0.001,1.001,len=50)
xP <- as.matrix(expand.grid("x1"=x_grid,"x2"=x_grid))
Emulate over this grid of prediction points and extract the expectation and variance:
# Evaluate emulator over 50x50=2500 prediction points xP
em_out <- t(apply(xP,1,simple_BL_emulator_v2,xD=xD,D=D,theta=0.25,sigma=250,E_f=350))
E_D_fx_mat <- matrix(em_out[,"ExpD_f(x)"],nrow=length(x_grid),ncol=length(x_grid))
Var_D_fx_mat <- matrix(em_out[,"VarD_f(x)"],nrow=length(x_grid),ncol=length(x_grid))
Here is the plotting function for our output:
### define filled contour plot function for emulator output ###
emul_fill_cont <- function(
cont_mat, # matrix of values we want contour plot of
cont_levs=NULL, # contour levels (NULL: automatic selection)
nlev=20, # approx no. of contour levels for auto select
plot_xD=TRUE, # plot the design runs TRUE or FALSE
xD=NULL, # the design points if needed
xD_col="green", # colour of design runs
x_grid, # grid edge locations that define xP
... # extra arguments passed to filled.contour
){
### Define contour levels if necessary ###
if(is.null(cont_levs)) cont_levs <- pretty(cont_mat,n=nlev)
### create the filled contour plot ###
filled.contour(x_grid,x_grid,cont_mat,levels=cont_levs,xlab="x1",ylab="x2",...,
plot.axes={axis(1);axis(2) # sets up plotting in contour box
contour(x_grid,x_grid,cont_mat,add=TRUE,levels=cont_levs,lwd=0.8) # plot contour lines
if(plot_xD) points(xD,pch=21,col=1,bg=xD_col,cex=1.5)}) # plot design points
}
Colour schemes:
library(viridisLite)
exp_cols <- plasma
var_cols <- function(n) hcl.colors(n, "blues3", rev = TRUE)
diag_cols <- turbo
Plot the emulator adjusted expectation and variance:
emul_fill_cont(cont_mat=E_D_fx_mat,cont_levs=seq(-50,1000,50),xD=xD,x_grid=x_grid,
color.palette=exp_cols, # this sets the colour scheme
main="Emulator Adjusted Expectation E_D[f(x)]")
emul_fill_cont(cont_mat=Var_D_fx_mat,cont_levs=NULL,xD=xD,x_grid=x_grid,
color.palette=var_cols,
main="Emulator Adjusted Variance Var_D[f(x)]")
In this instance we can plot the true 2-dimensional output of our SIRS model:
f <- NULL
for(i in 1:nrow(xP)){
parms = c( xP[i,1]*0.7+0.1, xP[i,2]*0.5, x3 = 0.04)
out <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000)) [101,"I"]
f <- c(f,out)
}
fxP_mat <- matrix(f,nrow=length(x_grid),ncol=length(x_grid))
emul_fill_cont(cont_mat=fxP_mat,cont_levs=seq(-50,1000,50),xD=xD,x_grid=x_grid,
color.palette=exp_cols,
main="True Computer Model Function f(x)" )
Plotting the diagnostics:
S_diag_mat <- (E_D_fx_mat - fxP_mat) / sqrt(Var_D_fx_mat)
emul_fill_cont(cont_mat=S_diag_mat,cont_levs=seq(-3,3,0.25),xD=xD,x_grid=x_grid,
xD_col="purple",
color.palette=diag_cols,
main="Emulator Diagnostics S_D[f(x)]")
We can see that the diagnostics look fine!
Derivative information can be extracted from the following adjoint SIRSS model:
### Define the SIRS Disease model. Note x1 = alpha_SI, x2 = alpha_IR, x3 = alpha_SR
### S = number of susceptible people, I = number of infected people, R = number of recovered people.
SIRS_Disease_Model_adjoint <- function(t, y, parms) {
with(as.list(parms),{
N <- y["S"] + y["I"] + y["R"]
dS = x3 * y["R"] - x1 * y["S"] * y["I"] / N
dI = x1 * y["S"] * y["I"] / N - x2 * y["I"]
dR = x2 * y["I"] - x3 * y["R"]
dSdx1 = x3*y["dRdx1"] -x1*y["dSdx1"]*y["I"]/N - x1*y["S"]*y["dIdx1"]/N - y["S"]*y["I"]/N
dSdx2 = x3*y["dRdx2"] -x1*y["dSdx2"]*y["I"]/N - x1*y["S"]*y["dIdx2"]/N
dSdx3 = x3*y["dRdx3"]+ y["R"] -x1*y["dSdx3"]*y["I"] /N - x1*y["S"]*y["dIdx3"]/N
dIdx1 = x1*y["dSdx1"]*y["I"]/N +x1*y["S"]*y["dIdx1"]/N +y["S"]*y["I"]/N -x2*y["dIdx1"]
dIdx2 = x1*y["dSdx2"]*y["I"]/N + x1*y["S"]*y["dIdx2"]/N -x2*y["dIdx2"] - y["I"]
dIdx3 = x1*y["dSdx3"]*y["I"] /N + x1*y["S"]*y["dIdx3"]/N - x2*y["dIdx3"]
dRdx1 = x2*y["dIdx1"] - x3*y["dRdx1"]
dRdx2 = y["I"] + x2*y["dIdx2"] -x3*y["dRdx2"]
dRdx3 = x2*y["dIdx3"] - y["R"] -x3*y["dRdx3"]
res <- c(dS, dI, dR, dSdx1,dSdx2,dSdx3,dIdx1,dIdx2,dIdx3,dRdx1,dRdx2,dRdx3)
list(res)
})
}
Bayes linear emulator using partial derivative information in both the \(x_1\) and \(x_2\) directions:
simple_BL_emulator_v2_dev<- function(x, # the emulator prediction point
xD, # the run input locations xD
D, # the run outputs D = (f(x^1),...,f(x^n))
theta = 1, # the correlation lengths
sigma = 1, # the prior SD sigma sqrt(Var[f(x)])
E_f=0, # prior expectation of f: E(f(x)) = 0
n=16, # # the number of design runs
n_x1=16, # the number of x1 derivatives
n_x2=16 # the number of x2 derivatives
){
# Define Covariance structure of f(x): Cov[f(x),f(xdash)]
Cov_fx_fxdash <- function(x,xdash) sigma^2 * exp(-sum((x-xdash)^2)/theta^2)
# Derivatives here are w.r.t x1
# Define Covariance structure of f(x): Cov[f'(x),f(xdash)]
Cov_fx_fxdash_dx1 <- function(x,xdash) -2*sigma^2 *(x[1]-xdash[1]) *exp(-sum((x-xdash)^2)/theta^2) /theta^2
# Define Covariance structure of f(x): Cov[f(x),f'(xdash)]
Cov_fx_fxdash_dx1dash <- function(x,xdash) 2*sigma^2 *(x[1]-xdash[1]) *exp(-sum((x-xdash)^2)/theta^2)/theta^2
# Define Covariance structure of f(x): Cov[f'(x),f'(xdash)]
Cov_fx_fxdash_dx1dx1dash <- function(x,xdash) -4*sigma^2 *(x[1]-xdash[1])^2 *exp(-sum((x-xdash)^2)/theta^2)/theta^4 + 2*sigma^2*exp(-sum((x-xdash)^2)/theta^2) /theta^2
# Derivatives here are w.r.t x2
# Define Covariance structure of f(x): Cov[f'(x),f(xdash)]
Cov_fx_fxdash_dx2 <- function(x,xdash) -2*sigma^2 *(x[2]-xdash[2]) *exp(-sum((x-xdash)^2)/theta^2) /theta^2
# Define Covariance structure of f(x): Cov[f(x),f'(xdash)]
Cov_fx_fxdash_dx2dash <- function(x,xdash) 2*sigma^2 *(x[2]-xdash[2]) *exp(-sum((x-xdash)^2)/theta^2)/theta^2
# Define Covariance structure of f(x): Cov[f'(x),f'(xdash)]
Cov_fx_fxdash_dx2dx2dash <- function(x,xdash) -4*sigma^2 *(x[2]-xdash[2])^2 *exp(-sum((x-xdash)^2)/theta^2)/theta^4 + 2*sigma^2*exp(-sum((x-xdash)^2)/theta^2) /theta^2
# Mixed partial derivatives
Cov_mixed <- function(x,xdash) -4*sigma^2 *(x[2]-xdash[2])*(x[1]-xdash[1]) *exp(-sum((x-xdash)^2)/theta^2)/theta^4
# Create E[D] vector
# Give the derivatives zero expectation
E_D <- c(rep(E_f,n),rep(0,n_x1),rep(0,n_x2))
#E_D <- c(rep(E_f,n+n_x1+n_x2))
# Create Var_D matrix
Var_D <- matrix(0,nrow=n+n_x1+n_x2,ncol=n+n_x1+n_x2)
# Keep this part of the matrix the same as in the non-derivative case
for(i in 1:n) for(j in 1:n) Var_D[i,j] <- Cov_fx_fxdash(xD[i,],xD[j,])
# Include the derivatives w.r.t x1
for(i in 1:n) for(j in (n+1):(n+n_x1)) Var_D[i,j] <- Cov_fx_fxdash_dx1dash(xD[i,],xD[j,])
for(i in (n+1):(n+n_x1)) for(j in 1:n) Var_D[i,j] <-Cov_fx_fxdash_dx1(xD[i,],xD[j,])
for(i in (n+1):(n+n_x1)) for(j in (n+1):(n+n_x1) ) Var_D[i,j] <-Cov_fx_fxdash_dx1dx1dash(xD[i,],xD[j,])
# Now include the derivatives w.r.t x2
for(i in 1:n) for(j in (n+n_x1+1):(n+n_x1+n_x2)) Var_D[i,j] <- Cov_fx_fxdash_dx2dash(xD[i,],xD[j,])
for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in 1:n) Var_D[i,j] <-Cov_fx_fxdash_dx2(xD[i,],xD[j,])
for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in (n+n_x1+1):(n+n_x1+n_x2) ) Var_D[i,j] <- Cov_fx_fxdash_dx2dx2dash(xD[i,],xD[j,])
for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in (n+1):(n+n_x1)) Var_D[i,j] <- Cov_mixed(xD[i,],xD[j,])
for(i in (n+1):(n+n_x1)) for(j in (n+n_x1+1):(n+n_x1+n_x2)) Var_D[i,j] <- Cov_mixed(xD[i,],xD[j,])
# Create E[f(x)]
E_fx <- E_f
# Create Var_f(x)
Var_fx <- sigma^2
# Create Cov_fx_D row vector
Cov_fx_D <- matrix(0,nrow=1,ncol=n+n_x1+n_x2)
# Covariance for our known runs
for(j in 1:n) Cov_fx_D[1,j] <- Cov_fx_fxdash(x,xD[j,])
# Covariance for x1 partial derivatives
for(j in (n+1):(n+n_x1)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx1dash(x,xD[j,])
# Covariance for x2 partial derivatives
for(j in (n+n_x1+1):(n+n_x1+n_x2)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx2dash(x,xD[j,])
#Perform Bayes Linear adjustment to find Adjusted Expectation and Variance of f(x)
ED_fx <- E_fx + Cov_fx_D %*% solve(Var_D) %*% (D - E_D)
VarD_fx <- Var_fx - Cov_fx_D %*% solve(Var_D) %*% t(Cov_fx_D)
### Return the emulator adjusted expectation and variance
return(c("ExpD_f(x)"=ED_fx,"VarD_f(x)"=VarD_fx))
}
Calculate all partial derivatives in the \(x_1\) and \(x_2\) direction:
# Starting configuration for the SIRSS model
ystart <- c(S = 850, I = 150, R = 0, dSdx1=0, dSdx2=0, dSdx3=0, dIdx1=0, dIdx2=0, dIdx3=0, dRdx1=0, dRdx2=0, dRdx3=0)
x1_dev <- NULL
for(i in 1:nrow(xD)){
parms = c(xD_scaled[i,1], xD_scaled[i,2], x3=0.04,dSdx1 = 0, dIdx1 = 0, dRdx1 = 0, dSdx2 = 0, dIdx2 = 0, dRdx2 = 0, dSdx3 = 0, dIdx3 = 0, dRdx3 = 0)
# Extract the output at t=10, this corresponds to the 101st row
infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model_adjoint, parms, maxsteps=200000)) [101,"dIdx1"]
x1_dev <- c(x1_dev,infected)
}
x1_dev <- x1_dev*0.7
# To get original scale just multiply by 0.7
x2_dev <- NULL
for(i in 1:nrow(xD_scaled)){
parms = c(xD_scaled[i,1], xD_scaled[i,2], x3=0.04, dSdx1 = 0, dIdx1 = 0, dRdx1 = 0, dSdx2 = 0, dIdx2 = 0, dRdx2 = 0, dSdx3 = 0, dIdx3 = 0, dRdx3 = 0)
# Extract the output at t=10, this corresponds to the 101st row
infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model_adjoint, parms, maxsteps=20000)) [101,"dIdx2"]
x2_dev <- c(x2_dev,infected)
}
x2_dev <- x2_dev*0.5
Emulate using this complete derivative information:
# Add this derivative information to D
D <- c(D,x1_dev,x2_dev)
# Record the input points for each partial derivative
xD <- rbind(xD,xD,xD)
em_out <- t(apply(xP,1,simple_BL_emulator_v2_dev,xD=xD,D=D,theta=0.25,sigma=250,E_f=350))
### store emulator output as matrices to aid plotting ###
E_D_fx_mat <- matrix(em_out[,"ExpD_f(x)"],nrow=length(x_grid),ncol=length(x_grid))
Var_D_fx_mat <- matrix(em_out[,"VarD_f(x)"],nrow=length(x_grid),ncol=length(x_grid))
emul_fill_cont(cont_mat=E_D_fx_mat,cont_levs=seq(-50,1000,50),xD=xD,x_grid=x_grid,
color.palette=exp_cols, # this sets the colour scheme
main="Emulator Adjusted Expectation E_D[f(x)]")
emul_fill_cont(cont_mat=Var_D_fx_mat,cont_levs=NULL,xD=xD,x_grid=x_grid,
color.palette=var_cols,
main="Emulator Adjusted Variance Var_D[f(x)]")
Plotting the diagnostics:
S_diag_mat <- (E_D_fx_mat - fxP_mat) / sqrt(Var_D_fx_mat)
emul_fill_cont(cont_mat=S_diag_mat,cont_levs=seq(-3,3,0.25),xD=xD,x_grid=x_grid,
xD_col="purple",
color.palette=diag_cols,
main="Emulator Diagnostics S_D[f(x)]")
Considering only the derivative information in the \(x_1\) direction:
simple_BL_emulator_v2_x1 <- function(x, # the emulator prediction point
xD, # the run input locations xD
D, # the run outputs D = (f(x^1),...,f(x^n))
theta = 1, # the correlation lengths
sigma = 1, # the prior SD sigma sqrt(Var[f(x)])
E_f = 0, # prior expectation of f: E(f(x)) = 0
n_x1 = 16 # the number of x1 derivatives
){
# store length of runs D
n <- 16
### Define Covariance structure of f(x): Cov[f(x),f(xdash)]
Cov_fx_fxdash <- function(x,xdash) sigma^2 * exp(-sum((x-xdash)^2)/theta^2)
# Derivatives are w.r.t x1:
# Define Covariance structure of f(x): Cov[f'(x),f(xdash)]
Cov_fx_fxdash_dx1 <- function(x,xdash) -2*sigma^2 *(x[1]-xdash[1]) *exp(-sum((x-xdash)^2)/theta^2) /theta^2
# Define Covariance structure of f(x): Cov[f(x),f'(xdash)]
Cov_fx_fxdash_dx1dash <- function(x,xdash) 2*sigma^2 *(x[1]-xdash[1]) *exp(-sum((x-xdash)^2)/theta^2)/theta^2
# Define Covariance structure of f(x): Cov[f'(x),f'(xdash)]
Cov_fx_fxdash_dx1dx1dash <- function(x,xdash) -4*sigma^2 *(x[1]-xdash[1])^2 *exp(-sum((x-xdash)^2)/theta^2)/theta^4 + 2*sigma^2*exp(-sum((x-xdash)^2)/theta^2) /theta^2
# Create E[D] vector
# Give the x1 partial derivatives zero expectation
E_D <- c(rep(E_f,n), rep(0,n_x1))
# Create Var_D matrix:
Var_D <- matrix(0,nrow=n+n_x1,ncol=n+n_x1)
# Keep this part of the matrix the same as emulating without derivative information
for(i in 1:n) for(j in 1:n) Var_D[i,j] <- Cov_fx_fxdash(xD[i,],xD[j,])
# Including the x1 partial derivatives
for(i in 1:n) for(j in (n+1):(n+n_x1)) Var_D[i,j] <- Cov_fx_fxdash_dx1dash(xD[i,],xD[j,])
for(j in 1:n) for(i in (n+1):(n+n_x1)) Var_D[i,j] <-Cov_fx_fxdash_dx1(xD[i,],xD[j,])
for(i in (n+1):(n+n_x1)) for(j in (n+1):(n+n_x1) ) Var_D[i,j] <-Cov_fx_fxdash_dx1dx1dash(xD[i,],xD[j,])
# Create E[f(x)]
E_fx <- E_f
# Create Var_f(x)
Var_fx <- sigma^2
# Create Cov_fx_D row vector
Cov_fx_D <- matrix(0,nrow=1,ncol=n+n_x1)
for(j in 1:n) Cov_fx_D[1,j] <- Cov_fx_fxdash(x,xD[j,])
# Include the x1 partial derivative information
for(j in (n+1):(n+n_x1)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx1dash(x,xD[j,])
#Perform Bayes Linear adjustment to find Adjusted Expectation and Variance of f(x)
ED_fx <- E_fx + Cov_fx_D %*% solve(Var_D) %*% (D - E_D)
VarD_fx <- Var_fx - Cov_fx_D %*% solve(Var_D) %*% t(Cov_fx_D)
# Return emulator adjusted expectation and variance
return(c("ExpD_f(x)"=ED_fx,"VarD_f(x)"=VarD_fx))
}
Emulate only considering the partial derivatives for all of the known runs in the \(x_1\) direction:
ystart <- c(S = 850, I = 150, R = 0)
# Perform 16 runs of the SIR model extracting the output for t=10 and store as D
D <- NULL
for(i in 1:nrow(xD_scaled)){
parms = c( xD_scaled[i,1], xD_scaled[i,2], x3 = 0.04)
infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000)) [101,"I"]
D <- c(D,infected)
}
### Define 50x50 grid of prediction points xP for emulator evaluation ###
x_grid <- seq(-0.001,1.001,len=50)
xP <- as.matrix(expand.grid("x1"=x_grid,"x2"=x_grid))
D <- c(D,x1_dev)
xD <- rbind(xD,xD)
em_out <- t(apply(xP,1,simple_BL_emulator_v2_x1,xD=xD,D=D,n_x1=16,theta=0.25,sigma=250,E_f=350))
### store emulator output as matrices to aid plotting ###
E_D_fx_mat <- matrix(em_out[,"ExpD_f(x)"],nrow=length(x_grid),ncol=length(x_grid))
Var_D_fx_mat <- matrix(em_out[,"VarD_f(x)"],nrow=length(x_grid),ncol=length(x_grid))
emul_fill_cont(cont_mat=E_D_fx_mat,cont_levs=seq(-50,1000,50),xD=xD,x_grid=x_grid,
color.palette=exp_cols, # this sets the colour scheme
main="Emulator Adjusted Expectation E_D[f(x)]")
emul_fill_cont(cont_mat=Var_D_fx_mat,cont_levs=NULL,xD=xD,x_grid=x_grid,
color.palette=var_cols,
main="Emulator Adjusted Variance Var_D[f(x)]")
Now just including all of the partial derivatives in the \(x_2\) direction:
simple_BL_emulator_v2_x2 <- function(x, # the emulator prediction point
xD, # the run input locations xD
D, # the run outputs D = (f(x^1),...,f(x^n))
theta = 1, # the correlation lengths
sigma = 1, # the prior SD sigma sqrt(Var[f(x)])
E_f = 0 , # prior expectation of f: E(f(x)) = 0
n_x2 = 16 # the number of x2 derivatives
){
# store length of runs D
n <- 16
# Define Covariance structure of f(x): Cov[f(x),f(xdash)]
Cov_fx_fxdash <- function(x,xdash) sigma^2 * exp(-sum((x-xdash)^2)/theta^2)
# Derivatives are w.r.t x2:
# Define Covariance structure of f(x): Cov[f'(x),f(xdash)]
Cov_fx_fxdash_dx2 <- function(x,xdash) -2*sigma^2 *(x[2]-xdash[2]) *exp(-sum((x-xdash)^2)/theta^2) /theta^2
# Define Covariance structure of f(x): Cov[f(x),f'(xdash)]
Cov_fx_fxdash_dx2dash <- function(x,xdash) 2*sigma^2 *(x[2]-xdash[2]) *exp(-sum((x-xdash)^2)/theta^2)/theta^2
# Define Covariance structure of f(x): Cov[f'(x),f'(xdash)]
Cov_fx_fxdash_dx2dx2dash <- function(x,xdash) -4*sigma^2 *(x[2]-xdash[2])^2 *exp(-sum((x-xdash)^2)/theta^2)/theta^4 + 2*sigma^2*exp(-sum((x-xdash)^2)/theta^2) /theta^2
# Create E[D] vector
# Give the x2 partial derivatives zero expectation
E_D <- c(rep(E_f,n),rep(0,n_x2))
# Create Var_D matrix:
Var_D <- matrix(0,nrow=n+n_x2,ncol=n+n_x2)
# Keep this part of the matrix the same as emulating without derivative information
for(i in 1:n) for(j in 1:n) Var_D[i,j] <- Cov_fx_fxdash(xD[i,],xD[j,])
# Now include the x2 partial derivatives
for(i in 1:n) for(j in (n+1):(n+n_x2)) Var_D[i,j] <- Cov_fx_fxdash_dx2dash(xD[i,],xD[j,])
for(j in 1:n) for(i in (n+1):(n+n_x2)) Var_D[i,j] <-Cov_fx_fxdash_dx2(xD[i,],xD[j,])
for(i in (n+1):(n+n_x2)) for(j in (n+1):(n+n_x2) ) Var_D[i,j] <-Cov_fx_fxdash_dx2dx2dash(xD[i,],xD[j,])
# Create E[f(x)]
E_fx <- E_f
# Create Var_f(x)
Var_fx <- sigma^2
# Create Cov_fx_D row vector
Cov_fx_D <- matrix(0,nrow=1,ncol=n+n_x2)
for(j in 1:n) Cov_fx_D[1,j] <- Cov_fx_fxdash(x,xD[j,])
for(j in (n+1):(n+n_x2)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx2dash(x,xD[j,])
# Perform Bayes Linear adjustment to find Adjusted Expectation and Variance of f(x)
ED_fx <- E_fx + Cov_fx_D %*% solve(Var_D) %*% (D - E_D)
VarD_fx <- Var_fx - Cov_fx_D %*% solve(Var_D) %*% t(Cov_fx_D)
# Return emulator adjusted expectation and variance
return(c("ExpD_f(x)"=ED_fx,"VarD_f(x)"=VarD_fx))
}
Rebuilding the emulator:
xD <- x_lhd
# Perform 16 runs of the SIR model extracting the output for t=10 and store as D
D <- NULL
for(i in 1:nrow(xD_scaled)){
parms = c( xD_scaled[i,1], xD_scaled[i,2], x3 = 0.04)
infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000)) [101,"I"]
D <- c(D,infected)
}
# Define 50x50 grid of prediction points xP for emulator evaluation
x_grid <- seq(-0.001,1.001,len=50)
xP <- as.matrix(expand.grid("x1"=x_grid,"x2"=x_grid))
D <- c(D,x2_dev)
xD <- rbind(xD,xD)
em_out <- t(apply(xP,1,simple_BL_emulator_v2_x2,xD=xD,D=D,theta=0.25,sigma=250,E_f=350))
# Store the emulator output as matrices
E_D_fx_mat <- matrix(em_out[,"ExpD_f(x)"],nrow=length(x_grid),ncol=length(x_grid))
Var_D_fx_mat <- matrix(em_out[,"VarD_f(x)"],nrow=length(x_grid),ncol=length(x_grid))
emul_fill_cont(cont_mat=E_D_fx_mat,cont_levs=seq(-50,1000,50),xD=xD,x_grid=x_grid,
color.palette=exp_cols, # this sets the colour scheme
main="Emulator Adjusted Expectation E_D[f(x)]")
emul_fill_cont(cont_mat=Var_D_fx_mat,cont_levs=NULL,xD=xD,x_grid=x_grid,
color.palette=var_cols,
main="Emulator Adjusted Variance Var_D[f(x)]")
Now emulate over the 2-dimensional input space of \(x_1\) and \(x_3\) at time t=10, keeping \(x_2=0.7\). Here \(x_1 \in [0.1,0.8]\) and \(x_3 \in [0,0.05]\) so we scale these input ranges to [0,1] for our emulation.
set.seed(29)
x_lhd <- lhd_maximin(8)
# Define run locations as the Maximin LHD
xD <- x_lhd
xD_scaled <- cbind("x2"=rep(0,8),"x3"=rep(0,8))
xD_scaled[,1] <- xD[,1]*0.5
xD_scaled[,2] <- xD[,2]*0.05
# Perform 16 runs of the SIR model extracting the output for t=10 and store as D
D <- NULL
for(i in 1:nrow(xD_scaled)){
parms = c( x1=0.7, xD_scaled[i,1], xD_scaled[i,2])
# Extract the output at t=10, this corresponds to the 101st row
infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000)) [101,"R"]
D <- c(D,infected)
}
Emulator which includes derivative information and different correlation length in each direction:
simple_BL_emulator_dev_theta<- function(x, # the emulator prediction point
xD, # the run input locations xD
D, # the run outputs D = (f(x^1),...,f(x^n))
theta = c(1,1), # the correlation lengths
sigma = 1, # the prior SD sigma sqrt(Var[f(x)])
E_f=0, # prior expectation of f: E(f(x)) = 0
n=8, # # the number of design runs
n_x1=8, # the number of x1 derivatives
n_x2=8 # the number of x2 derivatives
){
### Define Covariance structure of f(x): Cov[f(x),f(xdash)] ###
Cov_fx_fxdash <- function(x,xdash) sigma^2 * exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)
## Derivatives w.r.t x1 ##
### Define Covariance structure of f(x): Cov[f'(x),f(xdash)] ###
Cov_fx_fxdash_dx1 <- function(x,xdash) -2*sigma^2 *(x[1]-xdash[1]) *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2) /theta[1]^2
### Define Covariance structure of f(x): Cov[f(x),f'(xdash)] ###
Cov_fx_fxdash_dx1dash <- function(x,xdash) 2*sigma^2 *(x[1]-xdash[1]) *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)/theta[1]^2
### Define Covariance structure of f(x): Cov[f'(x),f'(xdash)] ###
Cov_fx_fxdash_dx1dx1dash <- function(x,xdash) -4*sigma^2 *(x[1]-xdash[1])^2 *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)/theta[1]^4 + 2*sigma^2*exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2) /theta[1]^2
## Derivatives w.r.t x2 ##
### Define Covariance structure of f(x): Cov[f'(x),f(xdash)] ###
Cov_fx_fxdash_dx2 <- function(x,xdash) -2*sigma^2 *(x[2]-xdash[2]) *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2) /theta[2]^2
### Define Covariance structure of f(x): Cov[f(x),f'(xdash)] ###
Cov_fx_fxdash_dx2dash <- function(x,xdash) 2*sigma^2 *(x[2]-xdash[2]) *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)/theta[2]^2
### Define Covariance structure of f(x): Cov[f'(x),f'(xdash)] ###
Cov_fx_fxdash_dx2dx2dash <- function(x,xdash) -4*sigma^2 *(x[2]-xdash[2])^2 *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)/theta[2]^4 + 2*sigma^2*exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2) /theta[2]^2
# Mixed partial derivatives
Cov_mixed <- function(x,xdash) -4*sigma^2 *(x[2]-xdash[2])*(x[1]-xdash[1]) *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)/(theta[1]^2*theta[2]^2)
# Create E[D] vector
E_D <- c(rep(E_f,n),rep(0,n_x1),rep(0,n_x2))
#E_D <- c(rep(E_f,n+n_x1+n_x2))
# Create Var_D matrix:
Var_D <- matrix(0,nrow=n+n_x1+n_x2,ncol=n+n_x1+n_x2)
# Keep this part of the matrix the same
for(i in 1:n) for(j in 1:n) Var_D[i,j] <- Cov_fx_fxdash(xD[i,],xD[j,])
# Include the derivatives w.r.t x1
for(i in 1:n) for(j in (n+1):(n+n_x1)) Var_D[i,j] <- Cov_fx_fxdash_dx1dash(xD[i,],xD[j,])
for(i in (n+1):(n+n_x1)) for(j in 1:n) Var_D[i,j] <-Cov_fx_fxdash_dx1(xD[i,],xD[j,])
for(i in (n+1):(n+n_x1)) for(j in (n+1):(n+n_x1) ) Var_D[i,j] <-Cov_fx_fxdash_dx1dx1dash(xD[i,],xD[j,])
# Now include the derivatives w.r.t x2
for(i in 1:n) for(j in (n+n_x1+1):(n+n_x1+n_x2)) Var_D[i,j] <- Cov_fx_fxdash_dx2dash(xD[i,],xD[j,])
for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in 1:n) Var_D[i,j] <-Cov_fx_fxdash_dx2(xD[i,],xD[j,])
for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in (n+n_x1+1):(n+n_x1+n_x2) ) Var_D[i,j] <- Cov_fx_fxdash_dx2dx2dash(xD[i,],xD[j,])
for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in (n+1):(n+n_x1)) Var_D[i,j] <- Cov_mixed(xD[i,],xD[j,])
for(i in (n+1):(n+n_x1)) for(j in (n+n_x1+1):(n+n_x1+n_x2)) Var_D[i,j] <- Cov_mixed(xD[i,],xD[j,])
# Create E[f(x)]
E_fx <- E_f
# Create Var_f(x)
Var_fx <- sigma^2
# Create Cov_fx_D row vector
Cov_fx_D <- matrix(0,nrow=1,ncol=n+n_x1+n_x2)
for(j in 1:n) Cov_fx_D[1,j] <- Cov_fx_fxdash(x,xD[j,])
for(j in (n+1):(n+n_x1)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx1dash(x,xD[j,])
for(j in (n+n_x1+1):(n+n_x1+n_x2)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx2dash(x,xD[j,])
### Perform Bayes Linear adjustment to find Adjusted Expectation and Variance of f(x) ###
ED_fx <- E_fx + Cov_fx_D %*% solve(Var_D) %*% (D - E_D)
VarD_fx <- Var_fx - Cov_fx_D %*% solve(Var_D) %*% t(Cov_fx_D)
### return emulator expectation and variance ###
return(c("ExpD_f(x)"=ED_fx,"VarD_f(x)"=VarD_fx))
}
Extract the derivative information from the adjoint model:
ystart <- c(S = 850, I = 150, R = 0, dSdx1=0, dSdx2=0, dSdx3=0, dIdx1=0, dIdx2=0, dIdx3=0, dRdx1=0, dRdx2=0, dRdx3=0)
x2_dev <- NULL
for(i in 1:nrow(xD)){
parms = c(x1=0.7,xD_scaled[i,1], xD_scaled[i,2],dSdx1 = 0, dIdx1 = 0, dRdx1 = 0, dSdx2 = 0, dIdx2 = 0, dRdx2 = 0, dSdx3 = 0, dIdx3 = 0, dRdx3 = 0)
# Extract the output at t=10, this corresponds to the 101st row
infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model_adjoint, parms, maxsteps=200000)) [101,"dRdx2"]
x2_dev <- c(x2_dev,infected)
}
x2_dev <- x2_dev*0.5
# To get original scale just multiply by 0.5
x3_dev <- NULL
for(i in 1:nrow(xD)){
parms = c(x1=0.7,xD_scaled[i,1], xD_scaled[i,2],dSdx1 = 0, dIdx1 = 0, dRdx1 = 0, dSdx2 = 0, dIdx2 = 0, dRdx2 = 0, dSdx3 = 0, dIdx3 = 0, dRdx3 = 0)
# Extract the output at t=10, this corresponds to the 101st row
infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model_adjoint, parms, maxsteps=200000)) [101,"dRdx3"]
x3_dev <- c(x3_dev,infected)
}
x3_dev <- x3_dev*0.05
# To get original scale just multiply by 0.5
Add in this derivative information and emulate using different correlation lenghts:
D <- c(D,x2_dev,x3_dev)
xD <- rbind(xD,xD,xD)
em_out <- t(apply(xP,1,simple_BL_emulator_dev_theta,xD=xD,D=D,theta=c(0.25,0.75),sigma=250,E_f=350))
E_D_fx_mat <- matrix(em_out[,"ExpD_f(x)"],nrow=length(x_grid),ncol=length(x_grid))
Var_D_fx_mat <- matrix(em_out[,"VarD_f(x)"],nrow=length(x_grid),ncol=length(x_grid))
emul_fill_cont(cont_mat=E_D_fx_mat,cont_levs=seq(0,700,50),xD=xD,x_grid=x_grid,
color.palette=exp_cols, # this sets the colour scheme
main="Emulator Adjusted Expectation E_D[f(x)]")
emul_fill_cont(cont_mat=Var_D_fx_mat,cont_levs=NULL,xD=xD,x_grid=x_grid,
color.palette=var_cols,
main="Emulator Adjusted Variance Var_D[f(x)]")